from unittest import result
import torch
from dataset.dataset import TPLinkerDataset

class TPLinkerDatasetBalancer:
    def __init__(self, dataset:TPLinkerDataset) -> None:
        self._dataset=dataset
        pass
    
    def get_weights4scaling(self):
        h2t_labels_statics = []
        h2h_labels_statics = []
        t2t_labels_statics = []

        for example in self._dataset:
            h2t_label = example["labels"]["h2t"]
            h2h_label = example["labels"]["h2h"]
            t2t_label = example["labels"]["t2t"]
            
            h2t_labels_statics.append(self.target_num_per_class(h2t_label, 2))
            h2h_labels_statics.append(self.target_num_per_class(h2h_label, 3))
            t2t_labels_statics.append(self.target_num_per_class(t2t_label, 3))
        
        h2t_labels_num = torch.stack(h2t_labels_statics,0).sum(0)
        h2h_labels_num = torch.stack(h2h_labels_statics,0).sum(0)
        t2t_labels_num = torch.stack(t2t_labels_statics,0).sum(0)
        
        total_h2t = h2t_labels_num.sum(0).item()
        total_h2h = h2h_labels_num.sum(0).item()
        total_t2t = t2t_labels_num.sum(0).item()
        
        h2t_scaling = (total_h2t - h2t_labels_num)/total_h2t
        h2h_scaling = (total_h2h - h2h_labels_num)/(total_h2h*2)
        t2t_scaling = (total_t2t - t2t_labels_num)/(total_t2t*2)
        
        return {
            "h2t":torch.tensor(h2t_scaling, dtype=torch.float).to("cpu").clone().detach().requires_grad_(False),
            "h2h":torch.tensor(h2h_scaling, dtype=torch.float).to("cpu").clone().detach().requires_grad_(False),
            "t2t":torch.tensor(t2t_scaling, dtype=torch.float).to("cpu").clone().detach().requires_grad_(False)
        }
    
    @staticmethod
    def target_num_per_class(label, class_num):
        
        
        cuda_condition = torch.cuda.is_available()
        device = torch.device("cuda:0" if cuda_condition else "cpu")
        
        results = torch.zeros((class_num,), dtype=torch.long).to(device)
        label = label.to(device)

        for i in range(class_num):
            results[i] += torch.sum(label == i).item()

        return results
    
